#ifndef CKKS_SRC_SCHEME_H_
#define CKKS_SRC_SCHEME_H_

#include <map>
#include <chrono>
#include <cstring>
#include <omp.h>
// #include <eigen3/Eigen/Dense>
// #include <eigen3/Eigen/Sparse>
#include "common.h"
#include "ciphertext.h"
#include "context.h"
#include "plaintext.h"
#include "secret_key.h"
#include "key.h"
#include "numb.h"


//static long ENCRYPTION = 0;
//static long MULTIPLICATION = 1;
//static long CONJUGATION = 2;

// using namespace Eigen;
using namespace std;

class Scheme {
private:
public:

    uint64_t ** polyTPools2;
    Key encKey;
    SwitchKey multKey;
    SwitchKey rotateKey;
    SwitchKey conjKey;
    Context &context;

    uint64_t *axKey;
    uint64_t *bxKey;

    uint64_t *polyPQPoolAx;
    uint64_t *polyPQPoolBx;
    uint64_t **polyTPools2Ax;
    uint64_t **polyTPools2Bx;

    // Scheme::LeftRotateFast && Scheme::Conjugate
    uint64_t *axax;
    uint64_t *bxbx;
    uint64_t *ntt_index;
    uint64_t *ax;
    uint64_t *axmult;
    uint64_t *bxmult;

    uint64_t *vx;
    uint64_t *ex;


    // Scheme::AddLeftRotKey
    uint64_t *sx;

    // Scheme::Encode,Scheme::EncodeSingle
    uint64_t *m;

    // Scheme::Decode
    complex<double> *res;

    // Scheme::MultAndEqual
    uint64_t *axaxi;
    uint64_t *axbx1;
    uint64_t *axbx2;

    // Scheme::Square
    uint64_t *axbx;

    // Scheme::MultByConst
    complex<double> *cnstvec;

    // Scheme::MultByConstVec
    uint64_t *mx_1;

    // Scheme::ConjugateAndEqual
    uint64_t *bxconj;
    uint64_t *bx_CAE;
    uint64_t *ax_CAE;


    // Scheme::ModUpToQAndEqual
//    uint64_t *ra;
//    uint64_t *rb;


    map<long, Key> keyMap;              // /< contain Encryption, Multiplication and Conjugation keys, if generated
    map<long, SwitchKey> leftRotKeyMap; // /< contain left rotation keys, if generated
    map<long, SwitchKey> switchKeyMap;

    vector<SwitchKey> switchKeyVector;


    // 构造函数
    Scheme(SecretKey &secretKey, Context &context);
    // 拷贝构造函数
    Scheme(const Scheme &other) = delete;
    Scheme &operator = (const Scheme &other) = delete;
    // 移动构造函数
    Scheme(Scheme &&other) = delete;
    Scheme &operator = (Scheme &&other) = delete;

    void MulScalarP(uint64_t *res, uint64_t *a);

    void MulScalaInt(Ciphertext &cipher, double scale);

    /* *
     * generates key for public encryption (key is stored in keyMap)
     */
    void AddEncKey(SecretKey &secretKey);

    /* *
     * generates key for multiplication (key is stored in keyMap)
     */
    void AddMultKey(SecretKey &secretKey);

    /* *
     * generates key for conjugation (key is stored in keyMap)
     */
    void AddConjKey(SecretKey &secretKey);

    /* *
     * generates key for left rotation (key is stored in leftRotKeyMap)
     */
    void AddLeftRotKey(SecretKey &secretKey, long rot);

    /* *
     * generates all keys for power-of-two left rotations (keys are stored in leftRotKeyMap)
     */
    void AddLeftRotKeys(SecretKey &secretKey);

    /* *
     * generates all keys for power-of-two right rotations (keys are stored in leftRotKeyMap)
     */
    void AddRightRotKeys(SecretKey &secretKey);

    Plaintext Encode(double *vals, long slots, long l);

    Plaintext Encode(complex<double> *vals, long slots, long l);

    Plaintext EncodeSingle(complex<double> val, long l);

    complex<double> *Decode(Plaintext &msg, __uint128_t scale);
    // complex<double>* Decode(Plaintext& msg, double scale);

    complex<double> DecodeSingle(Plaintext &msg);

    // Encryption (secret and public version)
    Ciphertext EncryptMsg(SecretKey &secretkey, Plaintext &message);

    Ciphertext EncryptMsg(Plaintext &message);

    Plaintext DecryptMsg(SecretKey &secretKey, Ciphertext &cipher);

    Ciphertext Encrypt(double *vals, long slots, long l);

    Ciphertext Encrypt(complex<double> *vals, long slots, long l);

    Ciphertext EncryptSk(SecretKey& secretKey, double* vals, long slots, long l);

    Ciphertext EncryptSk(SecretKey& secretKey, complex<double>* vals, long slots, long l);
    
    Ciphertext EncryptSingle(complex<double> val, long l);

    complex<double> *Decrypt(SecretKey &secretKey, Ciphertext &cipher);

    complex<double> DecryptSingle(SecretKey &secretKey, Ciphertext &cipher);

    // Homomorphic Negation
    Ciphertext Negate(Ciphertext &cipher);
    void NegateAndEqual(Ciphertext &cipher);

    // Homomorphic Addition
    Ciphertext Add(Ciphertext &cipher1, Ciphertext &cipher2);
    void Add(Ciphertext &res, Ciphertext &cipher1, Ciphertext &cipher2);
    void AddAndEqual(Ciphertext &cipher1, Ciphertext &cipher2);

    // Homomorphic Substraction
    Ciphertext Sub(Ciphertext &cipher1, Ciphertext &cipher2);

    void SubAndEqual(Ciphertext &cipher1, Ciphertext &cipher2);
    void Sub2AndEqual(Ciphertext &cipher1, Ciphertext &cipher2);

    // Homomorphic Multiplication
    Ciphertext Mult(Ciphertext &cipher1, Ciphertext &cipher2);
    void Mult(Ciphertext& res, Ciphertext &cipher1, Ciphertext &cipher2);

    void MultAndEqual(Ciphertext &cipher1, Ciphertext &cipher2);

    // void KeySwitch(int level, uint64_t *axaxKSW, std::vector<std::vector<std::vector<uint64_t>>>& ax,
    //                std::vector<std::vector<std::vector<uint64_t>>>& bx, uint64_t *axmultKSW, uint64_t *bxmultKSW);

    // uint64_t *ExternalProduct(int levelQ, std::vector<uint64_t *>& PolyTs,
    //     std::vector<std::vector<std::vector<uint64_t>>>& key);

    //modify
    // void KeySwitch(int level, uint64_t *axaxKSW, std::vector<std::vector<std::vector<uint64_t>>>& ax,
    //                std::vector<std::vector<std::vector<uint64_t>>>& bx, uint64_t *axmultKSW, uint64_t *bxmultKSW);

    // void ExternalProduct(int levelQ, std::vector<uint64_t *>& PolyTs,
    //     std::vector<std::vector<std::vector<uint64_t>>>& key, uint64_t *res);

    void KeySwitch(int level, uint64_t *axaxKSW, std::vector<std::vector<std::vector<uint64_t>>>& ax,
                   std::vector<std::vector<std::vector<uint64_t>>>& bx, uint64_t *axmultKSW, uint64_t *bxmultKSW);

    void ExternalProduct(int levelQ, std::vector<uint64_t *>& PolyTs,
        std::vector<std::vector<std::vector<uint64_t>>>& axkey, 
        std::vector<std::vector<std::vector<uint64_t>>>& bxkey,
        uint64_t *axres, uint64_t *bxres);

    // Homomorphic Squaring
    Ciphertext Square(Ciphertext &cipher);
    void Square(Ciphertext &res, Ciphertext &cipher);
    void SquareAndEqual(Ciphertext &cipher);

    // Homomorphic Operations with Constant
    Ciphertext MultByi(Ciphertext &cipher);
    Ciphertext DivByi(Ciphertext &cipher);

    Ciphertext MultPt(Ciphertext &cipher, Plaintext &plaintext);
    void MultPt(Ciphertext &res, Ciphertext &cipher, Plaintext &plaintext);
    void AddPt(Ciphertext &res, Ciphertext &cipher, Plaintext &plaintext);
    void AddPt(Ciphertext &cipher, Plaintext &plaintext);

    Ciphertext AddConst(Ciphertext &cipher, double cnst);
    Ciphertext AddConst(Ciphertext &cipher, complex<double> cnst);

    void AddConstAndEqual(Ciphertext &cipher, double cnst);

    Ciphertext MultByConst(Ciphertext &cipher, double cnst);
    Ciphertext MultByConst(Ciphertext &cipher, complex<double> cnst);

    Ciphertext MultByConstVec(Ciphertext &cipher, double *cnstVec, long slots);
    Ciphertext MultByConstVec(Ciphertext &cipher, complex<double> *cnstVec, long slots);

    void MultByConstAndEqual(Ciphertext &cipher, double cnst);

//    Ciphertext ReScaleBy(Ciphertext &cipher, long dl);
    void ReScaleByAndEqual(Ciphertext &cipher, long dl);

    void ReScaleToAndEqual(Ciphertext &cipher, long l);

    Ciphertext ModDownBy(Ciphertext &cipher, long dl);
    void ModDownByAndEqual(Ciphertext &cipher, long dl);

    Ciphertext ModDownTo(Ciphertext &cipher, long l);
    void ModDownToAndEqual(Ciphertext &cipher, long l);

    Ciphertext LeftRotateFast(Ciphertext &cipher, long rotSlots);
    void LeftRotateAndEqualFast(Ciphertext &cipher, long rotSlots);

    Ciphertext LeftRotateByPo2(Ciphertext &cipher, long logRotSlots);
    void LeftRotateByPo2AndEqual(Ciphertext &cipher, long logRotSlots);

    Ciphertext RightRotateByPo2(Ciphertext &cipher, long logRotSlots);
    void RightRotateByPo2AndEqual(Ciphertext &cipher, long logRotSlots);

    Ciphertext LeftRotate(Ciphertext &cipher, long rotSlots);
    void LeftRotateAndEqual(Ciphertext &cipher, long rotSlots);

    Ciphertext RightRotate(Ciphertext &cipher, long rotSlots);
    void RightRotateAndEqual(Ciphertext &cipher, long rotSlots);

    Ciphertext Conjugate(Ciphertext &cipher);
    void Conjugate(Ciphertext &res, Ciphertext &cipher);
    void ConjugateAndEqual(Ciphertext &cipher);

    void ModDownToQ0AndEqual(Ciphertext &cipher);

    void ModUpToQAndEqual(Ciphertext &cipher);

    Ciphertext ModUpToQ(Ciphertext &cipher);

    void SubSum(Ciphertext &cipher);
};

#endif
